import numpy as np
import time
import math
from bitarray import bitarray  # pip install bitarray
from typing import List

# 在文件顶部统一导入（避免重复 import）
import CompressMethodZlib
import CompressMethodGzip
import CompressMethodBz2
import CompressMethodRle
import CompressMethodHuffman
import CompressMethodZstd
import CompressMethodDelta


# 使用示例
compressors = [
    (0, "zlib", CompressMethodZlib.compressMethod, False),
    (1, "gzip", CompressMethodGzip.compressMethod, False),
    (2, "zstd", CompressMethodZstd.compressMethod, False),
    (3, "bz2", CompressMethodBz2.compressMethod, False),
    (4, "delta", CompressMethodDelta.compressMethod, False),
    (5, "rle", CompressMethodRle.compressMethod, True),
    (6, "huffman", CompressMethodHuffman.compressMethod, True),
]


def findMinCompressedLength(SnbsStr: str) -> list[bytes, int, int]:
    """
    对输入数据使用多种压缩算法，返回压缩后最小的字节长度。
    SnbsStr: 由01组成的字符串
    输出：best_compressedBytes, codeNo 算法编码, min_len 最小压缩后长度
    """
    print("    findMinCompressedLength函数调用...")
    t0 = time.perf_counter()
    min_len = len(SnbsStr)  # 初始值为原始长度（最坏情况：不压缩）
    bestCodeNo = -1
    bestMethod = "none"
    best_compressedBytes = None
    for codeNo, name, func, has_extra in compressors:
        if codeNo > 0:
            continue  # 加快测试速度
        try:
            result = func(SnbsStr)
            if has_extra:
                # RLE 和 Huffman 返回 (compressedBytes, _, _)
                compressedBytes = result[0]
            else:
                compressedBytes = result

            compLen = len(compressedBytes)
            print(f"    {name:10}: {compLen} bytes")  # 可选：打印各算法结果

            if compLen < min_len:
                min_len = compLen
                bestCodeNo = codeNo
                bestMethod = name
                best_compressedBytes = compressedBytes

        except Exception as e:
            print(f"    Error in {func}: {e}")
            continue  # 如果某个算法出错，跳过它

    print(f"    Best method: {bestCodeNo}, Min length: {min_len} bytes")
    print("    耗时:", time.perf_counter() - t0)
    print("    findMinCompressedLength函数调用结束")
    return best_compressedBytes, bestCodeNo, min_len


def GetSnbBitStrIswrong(seq: np.ndarray, sn: int) -> str:
    print("    getSnbStr函数调用...")
    t0 = time.perf_counter()
    if seq.dtype != np.float32:
        seq = seq.astype(np.float32)
    # 1. 先 reshape 成 (N,4) 再取最后一列
    raw_bytes = seq.view(np.uint8).reshape(-1, 4)[:, 3]  # shape (N,)
    # 2. 一次性压成位数组
    ba = bitarray(endian="big")
    ba.frombytes(raw_bytes.tobytes())  # 长度 N*8 位
    # 3. 提取每个字节的低 sn 位
    wanted = bitarray(endian="big")
    for i in range(len(seq)):
        start = i * 8 + (8 - sn)
        wanted.extend(ba[start : start + sn])

    print("    耗时:", time.perf_counter() - t0)
    print("    getSnbStr函数调用结束")
    return wanted.to01()


# 把所有的Snb取出来
def getSnbStrSlow(seq: list[np.float32], sn) -> str:
    print("    getSnbStr函数调用...")
    t0 = time.perf_counter()
    snbs = ""
    for i in range(len(seq)):
        num = seq[i].astype(np.float32)
        binaryRepresentation = num.view(np.uint32)
        # 将 uint32 转换为二进制字符串
        binaryStr = format(binaryRepresentation, "032b")  # 32 位二进制字符串
        snbs += binaryStr[-sn:]  # 取出
    print("    耗时:", time.perf_counter() - t0)
    print("    getSnbStr函数调用结束")
    return snbs


def GetSnbBitStr(seq: List[np.float32], sn: int) -> str:
    print("    GetSnbBitStr函数调用...")
    t0 = time.perf_counter()

    # 转为 numpy 数组
    arr = np.array(seq, dtype=np.float32)

    # 视为 uint32 获取 IEEE 754 位模式
    binary_view = arr.view(np.uint32)

    # 提取低 sn 位
    low_bits = binary_view & ((1 << sn) - 1)  # 得到 shape=(N,) 的 uint32 数组

    # ✅ 正确方式：使用 np.vectorize 向量化 format(x, 'b')
    @np.vectorize
    def to_binary_string(x):
        return format(x, 'b')  # 返回二进制字符串

    bin_strs = to_binary_string(low_bits)

    # 补前导零到 sn 位
    bin_strs_padded = np.char.zfill(bin_strs, sn)  # 左侧补零到长度 sn

    # 拼接成一个长字符串
    snbs = ''.join(bin_strs_padded)

    print("    耗时:", time.perf_counter() - t0)
    print("    GetSnbBitStr函数调用结束")
    return snbs

def EmbedWithEmbr(
    originalSequence: list[np.float32], EmbededBits: list[np.uint8], sn: int, m: int
) -> list[np.float32]:
    """
    将秘密信息嵌入到浮点数序列中
    参数:
        originalSequence: 原始浮点数序列, numpy.float32
        EmbededBits: 待嵌入的秘密信息，二进制列表, list[numpy.uint8], 元素为0或1
        sn: snb length,int
        m: modulo key,int
    返回: 嵌入秘密信息后的序列,是 非零元素的值数组, 它只包含矩阵中所有非零元素的值，是一个一维的 numpy.ndarray。
    """
    print("    EmbedWithEmbr函数调用...")
    t0 = time.perf_counter()
    # 复制原始序列
    seq = originalSequence.copy()
    # originalSequence为adata.X.data,是 非零元素的值数组, 它只包含矩阵中所有非零元素的值，是一个一维的 numpy.ndarray。
    seqLength = len(seq)
    Snbs = GetSnbBitStr(seq, sn)
    # with open('Snbs.txt', 'w', encoding='utf-8') as f: #将Snbs保存到txt文件中
    #     for item in Snbs:
    #         f.write(item)    
    assert len(Snbs) == seqLength * sn
    print(f"    压缩前送入的数据bit数: {len(Snbs)}")
    compressedBytes, codeNo, minBytesLen = findMinCompressedLength(Snbs)
    with open("log.txt", "a", encoding="utf-8") as f:
        f.write(f"\n压缩前送入的数据bit数: {len(Snbs)}")
        f.write(f"\n压缩后的byte长度{len(compressedBytes)}")
        # f.write(compressedBytes.hex() + "\n")  # 加换行符，方便后续查看

    print(f"    压缩后的数据字节数: {len(compressedBytes)}")
    print(f"    Method codeNo:{codeNo} Minimum compressed length:{minBytesLen} bytes")
    print("    打包头部...")
    bioLen = np.array(
        [math.ceil(len(EmbededBits) / 8)], dtype=np.uint32
    )  # 生物信息的字节
    bio_len_arr = np.array([bioLen], dtype=">u4")  # >u4 表示大端 uint32
    bio_len_uint8_list = bio_len_arr.view(np.uint8)  # 拆分为 4 个 uint8

    algorithm = np.array([codeNo], dtype=np.uint8)

    compLen = np.array([minBytesLen], dtype=np.uint32)  # 压缩后的字节长度
    compLen_arr = np.array([compLen], dtype=">u4")  # >u4 表示大端 uint32
    compLen_uint8_list = compLen_arr.view(np.uint8)  # 拆分为 4 个 uint8

    def bits_to_bytes(bits):
        """
        将 list[0,1] 转换为 list[int]（每个 int 是一个字节）.   不足 8 位时，末尾补 0
        """
        if not bits:
            return []
        # 补齐到 8 的倍数（末尾补 0）
        padding = (8 - len(bits) % 8) % 8
        padded_bits = bits + [0] * padding
        byte_list = []
        for i in range(0, len(padded_bits), 8):
            chunk = padded_bits[i : i + 8]
            # 将 8 位二进制转为一个字节（整数）
            byte_val = 0
            for bit in chunk:
                byte_val = (byte_val << 1) + bit
            byte_list.append(byte_val)
        return byte_list

    bioInfo = bits_to_bytes(EmbededBits)  # 生物信息
    # 转换为十进制
    compressedBytes = list(compressedBytes)

    annex_len = np.array([0], dtype=np.uint32)  # 附加 Annex 的字节长度
    annex_len_arr = np.array([annex_len], dtype=">u4")  # >u4 表示大端 uint32
    annex_len_uint8_list = annex_len_arr.view(np.uint8)  # 拆分为 4 个 uint8

    annexByts = np.array([], dtype=np.uint8)
    embedInfo = np.concatenate(
        (  # 所有元素的字节，所以bit数量应该是8的整数倍
            bio_len_uint8_list.flatten(),  # 4字节
            algorithm,  # 1字节
            compLen_uint8_list.flatten(),  # 4字节
            annex_len_uint8_list.flatten(),  # 4字节
            bioInfo,
            compressedBytes,  # 被压缩的部分，注意其他部分没有压缩，这一部分包括原始的所有snbs，因为是压缩，所以可能有填充的数据在里边
            annexByts,
        )
    )
    # print(embedInfo)
    int_array = embedInfo.astype(np.uint8)
    print(f"    原始序列长度：{seqLength} 嵌入的字节数：{len(int_array)}")
    # int_array = np.array([b for b in embedInfo.astype(bytes)], dtype=int)
    # bits = ""  # bits为01字符串
    # for b in int_array:
    #     # 将每个字节转为 8 位二进制字符串，如 '00000101'
    #     # print(f"    {b} type is {type(b)}")
    #     bin_str = format(b, "08b")  # 等价于 f"    {b:08b}"
    #     bits += bin_str
    bits = (
        np.unpackbits(int_array.astype(np.uint8), bitorder="big").astype(str).view("U1")
    )
    bits = "".join(bits)
    # 嵌入信息
    EmbededBits = bits  # 要嵌入的EmbededBits, EmbededBits每个元素为np.uint8，值为0或者1
    EmbededBitsLen = len(EmbededBits)
    assert EmbededBitsLen % 8 == 0
    # print(f"    嵌入的信息: {EmbededBits}")

    # 计算最大可嵌入比特数
    maxCapacity = seqLength * sn
    if EmbededBitsLen > maxCapacity:
        raise ValueError(
            f"    秘密信息过长，最大容量为{maxCapacity}位，当前为{len(EmbededBits) * sn}位"
        )
    print("    开始嵌入...")
    len1 = len(EmbededBits)
    remainder = len1 % sn
    EmbededBits += "0" * (sn - remainder)
    len1 = len(EmbededBits)
    total_sn = len(EmbededBits) // sn
    bits_array = np.array(list(EmbededBits), dtype=np.uint8).reshape(total_sn, sn)
    # print(bits_array[:20])
    snb_ints = bits_array @ (1 << np.arange(sn - 1, -1, -1))
    # print(snb_ints[:20])
    snb_ints = (snb_ints + m) % (2**sn)
    # print(snb_ints[:20])
    snb_ints = snb_ints.astype(np.uint32)
    # print(snb_ints[:20])
    uint32_view = seq[:total_sn].view(np.uint32)
    # print(uint32_view[:20])
    mask = np.uint32(0xFFFFFFFF) ^ np.uint32((1 << sn) - 1)
    # print(mask)  #4294967288   11111111111111111111111111111000  @ sn == 3
    uint32_view[:] = (uint32_view & mask) | snb_ints
    # print(uint32_view[:20])
    print("    耗时:", time.perf_counter() - t0)
    print("    EmbedWithEmbr函数调用结束")
    return seq


def ExtractWithEmbr(
    embeddedSequence: list[np.float32], sn: int, m: int
) -> list[bytes, bytes]:
    """
    从嵌入信息的浮点数序列中提取信息，解密后返回
    参数:
        embeddedSequence: 嵌入信息后的序列，numpy.float32
        sn: snb length,int
        m: modulo,int
    返回:
        bioStr字符串,解密后的bit list
    """
    print("    ExtractWithEmbr函数调用...")
    extracted = []
    seqLen = len(embeddedSequence)
    extracted = GetSnbBitStr(embeddedSequence, sn)
    # extracted1 = getSnbStrSlow(embeddedSequence, sn)
    # 需要解密
    n_groups = len(extracted) // sn
    int3 = np.fromiter(extracted, dtype=np.uint8).reshape(n_groups, sn)
    vals = int3 @ (1 << np.arange(sn-1, -1, -1))      # 形状 (n_groups,)
    # print(vals[:20])       
    vals = (vals - m + (1 << sn)) % (1 << sn) #模运算
    vals = np.asarray(vals, dtype=np.uint8)          # 强制无符号 8 位
    # print(vals[:20])       
    bits8 = np.unpackbits(vals, bitorder='big')   # 返回长度 len(vals)*8 的一维 01 数组
    # print(bits8[:20])          
    bits8_2d = bits8.reshape(-1,8)
    # print(bits8_2d[:20])          
    snb_2d = bits8_2d[:, -sn:]              # shape (len(vals), sn)
    # print(snb_2d[:20])          
    # snb_1d = ''.join(snb_2d.astype(str).ravel())  #很耗内存，每个元素从 1 字节的整数变成最多 2~10 字节的 Python 字符串对象。内存占用可能暴涨 10~50 倍！
     # 分块生成二进制字符串，避免内存爆炸
    # 改为以下代码：
    def bits_to_string_generator(snb_2d):
        for row in snb_2d:
            yield ''.join(row.astype(str))  # 每行转字符串，但不保留大数组

    snb_1d = ''.join(bits_to_string_generator(snb_2d))
    extracted = snb_1d
    # 提取信息
    assert len(extracted) == seqLen * sn
    # 4字节
    bio_len_str = extracted[0 : 4 * 8]
    bioLen = np.uint32(int(bio_len_str, 2))  # 所有byte的长度
    # 1字节
    algorithm_str = extracted[4 * 8 : 5 * 8]
    algorithm = np.uint8(int(algorithm_str, 2))
    # 4字节
    compressedLen_str = extracted[5 * 8 : 9 * 8]
    compressedLen = np.uint32(int(compressedLen_str, 2))
    # 4字节
    annex_len_str = extracted[9 * 8 : 13 * 8]
    annex_len = np.uint32(int(annex_len_str, 2))

    bioStrStart = 13 * 8
    bioStr = ""
    for i in range(bioLen):
        str1 = extracted[bioStrStart + i * 8 : bioStrStart + i * 8 + 8]
        bioStr += chr(np.uint8(int(str1, 2)))
    print(f"    biostr: {bioStr}")

    compressedStrStart = bioStrStart + bioLen * 8
    compressedUint8 = []
    for i in range(compressedLen):
        str1 = extracted[compressedStrStart + i * 8 : compressedStrStart + +i * 8 + 8]
        compressedUint8.append(np.uint8(int(str1, 2)))
    compressedBytes = bytes(compressedUint8)
    print(f"    提取的未解压的字节长度: {len(compressedBytes)}")


    match algorithm:
        case 0:  # 0: "zlib"
            print("    解压算法：zlib")
            restoredSnbBytes = CompressMethodZlib.decompressMethod(compressedBytes)
        case 1:  # 1: "gzip"
            print("    解压算法：gzip")
            restoredSnbBytes = CompressMethodGzip.decompressMethod(compressedBytes)
        case 2:  # 2: "zstd"
            print("    解压算法：zstd")
            restoredSnbBytes = CompressMethodZstd.decompressMethod(compressedBytes)
        case 3:  # 3: "bz2"
            print("    解压算法：bz2")
            restoredSnbBytes = CompressMethodBz2.decompressMethod(compressedBytes)
        case 4:  # 4: "delta"
            print("    解压算法：delta")
            restoredSnbBytes = CompressMethodDelta.decompressMethod(compressedBytes)
        case 5:  # 5: "rle"
            print("    解压算法：rle")
            restoredSnbBytes = CompressMethodRle.decompressMethod(compressedBytes)
        case 6:  # 6: "huffman"
            print("    解压算法：huffman")
            restoredSnbBytes = CompressMethodHuffman.decompressMethod(compressedBytes)
        case _:  # 相当于 default
            pass
    #restoredSnbBytes：返回的是一个一个的bit
    print(f"    解压后的bit数：{len(restoredSnbBytes)}")
    with open("log.txt", "a", encoding="utf-8") as f:
        f.write(f"\n解压前的byte长度{len(compressedBytes)}")
        f.write(f"\n解压后的bit数：{len(restoredSnbBytes)}")
        # f.write(compressedBytes.hex() + "\n")  # 加换行符，方便后续查看

    return bioStr, restoredSnbBytes  # 是bit list


# 示例用法
if __name__ == "__main__":
    pass
